#include "wingcoroutine_p.h"

#include <algorithm>

using namespace Wing;

#define DECLARE_STRUCT_OFFSET(type, member) [member] "i"(offsetof(type, member))

WingCoroutine_p::WingCoroutine_p(size_t stackSize)
    : _stackSize(stackSize), _sm(0) {
  for (int i = 1; i <= 10; ++i) {
    _ids.push_back(i);
  }
  _idLimit = 10;
}

void WingCoroutine_p::swapContextCaller(Wing::WingCoroutine *p) {
  bool found = false;
  while (true) {
    while (!this->_apcs.empty()) {
      auto co_apc = this->_apcs.front();
      this->_apcs.pop_front();
      co_apc();
    }

    if (_threads.size() < 2) {
      break;
    }
    auto th = _currentThread;
    th++;
    if (th == _threads.end()) {
      th = _threads.begin();
    }

    switch (th->state) {
    case ThreadState::READY: {
      auto oldth = _currentThread;
      auto newth = th;
      _currentThread = th;
      swapContext(*oldth, *newth, p);
      found = true;
    } break;
    case Wing::ThreadState::DEAD: {
      ::operator delete[](th->stack);
      if (th->param) {
        ::operator delete(th->param);
      }
      requeueID(th->id);
      _threads.erase(th);
      continue;
    } break;
    default:
      break;
    }

    if (found) {
      break;
    }
  }
}

bool WingCoroutine_p::sendCOAPC(const Wing::CO_APC &co_apc) {
  if (this->_hostThread.get_id() == std::thread::id()) {
    return false;
  }
  this->_apcs.push_back(co_apc);
  return true;
}

size_t WingCoroutine_p::stackSize() const { return _stackSize; }

ThreadID WingCoroutine_p::getID() {
  ThreadID id = 0;
  if (_ids.size()) {
    id = _ids.front();
    _ids.pop_front();
  } else {
    id = _idLimit + 1;
    for (int i = _idLimit + 2; i <= _idLimit + 10; ++i) {
      _ids.push_back(i);
    }
    _idLimit += 10;
  }
  return id;
}

void WingCoroutine_p::requeueID(Wing::ThreadID ID) {
  if (ID <= _idLimit && std::find(_ids.begin(), _ids.end(), ID) == _ids.end()) {
    _ids.push_front(ID);
  }
}

void WingCoroutine_p::swapContext(UNUSED const Wing::ThreadInfo &oldth,
                                  UNUSED const Wing::ThreadInfo &newth,
                                  UNUSED Wing::WingCoroutine *p) {
  asm volatile("pushq %rax;"
               "pushq %rbx;"
               "pushq %rcx;"
               "pushq %rdx;"
               "pushq %rbp;"
               "pushq %rsi;"
               "pushq %r8;"
               "pushq %r9;"
               "pushq %r10;"
               "pushq %r11;"
               "pushq %r12;"
               "pushq %r13;"
               "pushq %r14;"
               "pushq %r15;"
               "pushfq;");

#if defined(OS_WINDOWS)
  // rcx: old, rdx: new, r8: p
  asm volatile("movq %%rsp,%c[stackPC](%%rcx);" ::DECLARE_STRUCT_OFFSET(
                   ThreadInfo, stackPC)
               : "rcx", "rdx", "r8");
  asm volatile("movq %0,%c[state](%%rcx);" ::"i"(ThreadState::READY),
               DECLARE_STRUCT_OFFSET(ThreadInfo, state)
               : "rcx", "rdx", "r8");

  asm volatile("movq %c[pc](%%rdx),%%rax;"
               "test %%rax,%%rax;"
               //> if not first started
               "jz pcnull%=;"
               // r10 : function handler
               "movq %%rax,%%r10;"
               "xor %%eax,%%eax;"
               "movq %%rax,%c[pc](%%rdx);"
               "movq %c[stackPC](%%rdx),%%rbp;"
               "movq %%rbp,%%rsp;"
               // fill up stack info
               "lea idle%=(%%rip), %%rax;"
               "pushq %%rax;"
               // now r11: new
               "movq %%rdx,%%r11;"
               "movq %%r8,%%rcx;"
               "movq %c[param](%%r11),%%rdx;"
               "callq *%%r10;"
               "pcnull%=:;"
               //> if started
               "movq %1,%c[state](%%rdx);"
               "movq %c[stackPC](%%rdx),%%rsp;"
               "jmp sw%=;"

               "idle%=:;"
               "call %P0;"
               "jmp idle%=;"
               "sw%=:;" ::"ic"(swapContext),
               "i"(ThreadState::RUNNING),
               DECLARE_STRUCT_OFFSET(ThreadInfo, state),
               DECLARE_STRUCT_OFFSET(ThreadInfo, pc),
               DECLARE_STRUCT_OFFSET(ThreadInfo, param),
               DECLARE_STRUCT_OFFSET(ThreadInfo, stack),
               DECLARE_STRUCT_OFFSET(ThreadInfo, stackPC)
               : "rax", "rcx", "rdx", "r8");
#elif defined(OS_LINUX)
  // rdi: old, rsi: new, rdx: p
  asm volatile("movq %%rsp,%c[stackPC](%%rdi);" ::DECLARE_STRUCT_OFFSET(
                   ThreadInfo, stackPC)
               : "rdi", "rsi", "rdx");
  asm volatile("movq %0,%c[state](%%rdi);" ::"i"(ThreadState::SLEEP),
               DECLARE_STRUCT_OFFSET(ThreadInfo, state)
               : "rdi", "rsi", "rdx");

  asm volatile("movq %c[pc](%%rsi),%%rax;"
               "test %%rax,%%rax;"
               //> if not first started
               "jz pcnull%=;"
               // rbx : function handler
               "movq %%rax,%%rbx;"
               "xor %%eax,%%eax;"
               "movq %%rax,%c[pc](%%rsi);"
               "movq %c[stackPC](%%rsi),%%rbp;"
               "movq %%rbp,%%rsp;"
               // fill up stack info
               "lea idle%=(%%rip), %%rax;"
               "pushq %%rax;"
               // now r11: new
               "movq %%rsi,%%r11;"
               "movq %%rdx,%%rdi;"
               "movq %c[param](%%r11),%%rsi;"
               "callq *%%rbx;"
               "pcnull%=:;"
               //> if started
               "movq %1,%c[state](%%rsi);"
               "movq %c[stackPC](%%rsi),%%rsp;"
               "jmp sw%=;"

               "idle%=:;"
               "call %P0;"
               "jmp idle%=;"
               "sw%=:;" ::"ic"(swapContext),
               "i"(ThreadState::RUNNING),
               DECLARE_STRUCT_OFFSET(ThreadInfo, state),
               DECLARE_STRUCT_OFFSET(ThreadInfo, pc),
               DECLARE_STRUCT_OFFSET(ThreadInfo, param),
               DECLARE_STRUCT_OFFSET(ThreadInfo, stack),
               DECLARE_STRUCT_OFFSET(ThreadInfo, stackPC)
               : "rax", "rdi", "rsi", "rdx");
#else
#error "Unsuported operating system!"
#endif

  asm volatile("popfq;"
               "popq %r15;"
               "popq %r14;"
               "popq %r13;"
               "popq %r12;"
               "popq %r11;"
               "popq %r10;"
               "popq %r9;"
               "popq %r8;"
               "popq %rsi;"
               "popq %rbp;"
               "popq %rdx;"
               "popq %rcx;"
               "popq %rbx;"
               "popq %rax;"
               "ret;");
}
